import java.io.*;
import java.security.*;
import java.security.spec.*;
import java.util.Random;

import javax.crypto.*;
import javax.crypto.spec.*;

import org.apache.commons.codec.binary.Base64;

import com.pernix.log.ConsoleLogging;
import com.pernix.log.ILog;

/*
 * 	Name 		::: Akawit Tangkupatanonth
 * 	Date 		::: 21/03/2013
 * 	Description ::: Class for encryption message with AES256, Key Exchange with RSA1024, integrity with SHA256.
 */

public class EncryptAndDecrypt {
	
	//================================================================================
    // Properties
    //================================================================================

	private ILog logger;
	private Cipher cipherAES;
	private Cipher cipherRSA;
	private PublicKey publicKeyRSA;
	private PublicKey interlocutorPublicKey;
	private PrivateKey privateKeyRSA;
    private KeyGenerator kgen;
    private KeyPairGenerator keyGen;
    private SecretKeySpec skeySpec;
    private byte[] getinterlocutorPublicKey = null;
    private byte[] rawKey = null;
    private byte[] decryptedRawKey = null;
    private byte[] encryptedAesKey = null;
    private byte[] integrityString = null;
    private byte[] encryptedHMACKey = null;
    private FileOutputStream fospb = null;
    private FileOutputStream fospv = null;
    private FileInputStream fispb = null;
    private FileInputStream fispv = null;
    private File fspb = null;
    private File fspv = null;
    private ByteArrayOutputStream baospb = null;
    private ByteArrayOutputStream baospv = null;
    private X509EncodedKeySpec keySpec = null;
    private PKCS8EncodedKeySpec keySpecpv = null;
    private KeyFactory keyFactory = null;
    private String myIntegrityMessage = "";
    private String HMACKey = "";
    private IvParameterSpec ivspec = null;
    
    static final String AB = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ";
    static Random rnd = new Random();
    
    //================================================================================
    // Constructors
    //================================================================================
    
    public EncryptAndDecrypt(ILog logger) {
        try {
        	this.logger = logger;
        	keyGen = KeyPairGenerator.getInstance("RSA");
        	keyGen.initialize(1024);
        	
            cipherAES = Cipher.getInstance("AES/CBC/PKCS5Padding");
            cipherRSA = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            
            kgen = KeyGenerator.getInstance("AES");
            kgen.init(256);
            
            byte[] iv = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
            ivspec = new IvParameterSpec(iv);
            
        } catch (Exception ex) {
        	ex.printStackTrace();
        	logger.Error("Error ::: " + ex.getMessage());
        } 
    }
    
    //================================================================================
    // Method
    //================================================================================
    
    // แปลง Array of Byte ให้อยู่ในรูปข้อความปกติ
    public String asHex (byte buf[]) {
    	
    	return Base64.encodeBase64String(buf);
  	}
    
    // แปลงข้อความ ให้อยู่ในรูป Array of Byte
    public byte[] deHex (String text) {
    	
    	return Base64.decodeBase64(text);
  	}
    
    // แปลงข้อความให้พร้อมส่ง
    public String prePareEncryptMessage(String Message, String whoReceive) throws Exception {
    	logger.Info("========== Start Prepare to encrypt message before send ========== ");
    	logger.Info("Send to                                 ::: " + whoReceive);
    	logger.Info("Message original                        ::: " + Message);
    	
    	myIntegrityMessage = "message=" + Message;
    	
    	generateKeyAES(); // Generate AES Key
    	integrityString(randomString(), myIntegrityMessage); // Generate HMAC Key And Integrity Message
    	
    	logger.Info("Message sended                          ::: " + encryptMessage(Message)+"&"+asHex(encryptedAesKey)+"&"+asHex(encryptedHMACKey)+"&"+asHex(integrityString));
    	logger.Info("========== End Prepare to encrypt message before send   ========== ");
    	
    	return encryptMessage(Message)+"&"+asHex(encryptedAesKey)+"&"+asHex(encryptedHMACKey)+"&"+asHex(integrityString);
    }
    
    // แปลงข้อความที่ได้รับ
    public String prePareDecryptMessage(String FullMessage, String username) throws Exception {
    	
    	String returnMessage = "";
    	
    	logger.Info("========== Start Prepare to decrypt message before show  ========== ");
    	logger.Info("Receiver                                ::: " + username);
    	logger.Info("EncryptMessage                          ::: " + FullMessage);
    	
    	String messageRecieve = FullMessage.split("&")[0]; // Message
    	SecretKeySpec SecretKeyAESReceive = decryptKeyAES(Base64.decodeBase64(FullMessage.split("&")[1])); // AESKey
    	String IntegrityKeyReceive = decryptKeyIntegrity(Base64.decodeBase64(FullMessage.split("&")[2])); // HMACKey
    	String IntegrityMessageReceive = new String(Base64.decodeBase64(FullMessage.split("&")[3])); // HMACMessage
    	
    	logger.Info("messageRecieve                          ::: " + messageRecieve);
    	logger.Info("SecretKeyAESReceive                     ::: " + SecretKeyAESReceive);
    	logger.Info("IntegrityKeyReceive                     ::: " + IntegrityKeyReceive);
    	logger.Info("IntegrityMessageReceive                 ::: " + IntegrityMessageReceive);
    	
    	String IntegrityMessage = "message=" + decrypt(messageRecieve, SecretKeyAESReceive);
    	
    	if(compairIntegrity(IntegrityMessage, IntegrityKeyReceive, IntegrityMessageReceive)){
    		returnMessage = decrypt(messageRecieve, SecretKeyAESReceive);
    	}
		
    	logger.Info("========== End Prepare to decrypt message before show    ========== ");
    	
		return returnMessage;
    }
    
    // สุ่มข้อความเพื่อทำ Key HMAC
    private String randomString() {
    	StringBuilder sb = new StringBuilder( 5 );
    	for( int i = 0; i < 5; i++ ) 
    		sb.append( AB.charAt( rnd.nextInt(AB.length()) ) );
    	
    	logger.Info("Generate Random Key HMAC                ::: " + sb.toString());
    	return HMACKey = sb.toString();
	}
    
    // แปลงข้อความให้อยู่ในรูปแบบ Hash
    private void integrityString(String Key, String messageInte) {
		try{
			 logger.Info("Before convert message in hash format   ::: " + messageInte);
			
			 Mac sha256_HMAC = Mac.getInstance("HmacSHA256");
			 SecretKeySpec secret_key = new SecretKeySpec(Key.getBytes(), "HmacSHA256");
			 sha256_HMAC.init(secret_key);
			 logger.Info("Secret key for HMAC is                  ::: " + secret_key);
			 integrityString = sha256_HMAC.doFinal(messageInte.getBytes());
			 
			 logger.Info("After convert message in hash format    ::: " + integrityString);
			 
			 encryptKeyHMAC();
			
		} catch(Exception ex){
			throw new RuntimeException(ex);
		}
	}
    
    // เปรียบเทียบ HMAC
    private boolean compairIntegrity(String message, String Key, String messageInte) throws Exception {
    	Mac sha256_HMAC = Mac.getInstance("HmacSHA256");
    	SecretKeySpec secret_key = new SecretKeySpec(Key.getBytes(), "HmacSHA256");
    	sha256_HMAC.init(secret_key);
		 
    	logger.Info("Compare HMAC (HMAC Sender)              ::: " + messageInte);
    	logger.Info("Compare HMAC (HMAC Receiver)            ::: " + new String(sha256_HMAC.doFinal(message.getBytes())));
		 
    	if(new String(sha256_HMAC.doFinal(message.getBytes())).equals(messageInte)) return true;
    	else return false;
    }
    
    // ตรวจสอบ Public Key และ Private Key ก่อนทำการสร้าง
    public void checkAndgenerateKey() throws Exception {
        int theByte;
        fspb = new File("chatSecurepub");
        fspv = new File("chatSecureprv");
        
        
        
        // ตรวจสอบว่ามี Public Key และ Private หรือไม่ ถ้าไม่ให้ทำการสร้างขึ้นมาใหม่
        if(!fspb.exists() || !fspv.exists()) {
        	logger.Info("Check public key or private key already ::: false");
        	generateKeyRSA();
        }
        else{
        	try {
        		logger.Info("Check public key and private key already::: true");
        		
        		fispb = new FileInputStream("chatSecurepub");
        		fispv = new FileInputStream("chatSecureprv");
        		
        		baospb = new ByteArrayOutputStream();
        		baospv = new ByteArrayOutputStream();

        	    theByte = 0;
        	    while ((theByte = fispb.read()) != -1)
        	    {
        	    	baospb.write(theByte);
        	    }
        	    
        	    theByte = 0;
        	    while ((theByte = fispv.read()) != -1)
        	    {
        	    	baospv.write(theByte);
        	    }
        	    
        	    byte[] keyBytespublic = baospb.toByteArray();
        	    byte[] keyBytesprivate = baospv.toByteArray();
        	    
        	    keyFactory = KeyFactory.getInstance("RSA");
        	    
        	    keySpec = new X509EncodedKeySpec(keyBytespublic);
        	    publicKeyRSA = keyFactory.generatePublic(keySpec);
        	    
        	    logger.Info("Get public key from file                ::: " + publicKeyRSA);
        	    
        	    keySpecpv = new PKCS8EncodedKeySpec(keyBytesprivate);
        	    privateKeyRSA = keyFactory.generatePrivate(keySpecpv);
        	    
        	    logger.Info("Get private key from file               ::: " + privateKeyRSA);
        	    
        	}
        	catch(Exception ex){
        		logger.Info("Error ::: " + ex.getMessage());
        		generateKeyRSA();
        	}
        	finally
        	{
        		fispb.close();
        	    fispv.close();
        	    baospb.close();
        	    baospv.close();
        	}
        }
    }
    
    // สร้าง Public Key และ Private Key
    private void generateKeyRSA() throws Exception {
        KeyPair keyPair = null;
        
        try {
        	KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
        	fospb = new FileOutputStream("chatSecurepub");
            fospv = new FileOutputStream("chatSecureprv");
            
            keyPair = keyGen.generateKeyPair();
            publicKeyRSA = keyPair.getPublic();
            privateKeyRSA = keyPair.getPrivate();
            
        	fospb.write(publicKeyRSA.getEncoded());
        	fospv.write(privateKeyRSA.getEncoded());
        	
        	logger.Info("Generate public key and write in file   ::: " + publicKeyRSA);
        	logger.Info("Generate private key and write in file  ::: " + publicKeyRSA);

        } catch (NoSuchAlgorithmException ex) {
        	ex.printStackTrace();
        } finally {
        	fospb.close();
            fospv.close();
        }
    }
    
    // สร้าง Secret Key AES
    private void generateKeyAES() throws Exception {
        SecretKey aeskey = kgen.generateKey();
        rawKey = aeskey.getEncoded();
        
        logger.Info("Get Original key AES[SecretKey]         ::: " + aeskey);
    	logger.Info("Get Original key AES[Array]             ::: " + rawKey);
    	
        encryptKeyAES();
    }
    
    // นำข้อความเข้ารหัส AES
    private String encryptMessage(String message) throws Exception {
        String cypheredtext="";
        
        logger.Info("Before encrypt message with AES         ::: " + message);
        
        skeySpec = new SecretKeySpec(rawKey, "AES");
    	cipherAES.init(Cipher.ENCRYPT_MODE, skeySpec, ivspec);
        byte[] clearbyte = message.getBytes();
        
        byte[] cyphered = cipherAES.doFinal(clearbyte);
        cypheredtext = new String(Base64.encodeBase64(cyphered),"UTF8");
        
        logger.Info("After encrypt message with AES          ::: " + cypheredtext);

        return cypheredtext;
    }
    
    // ถอดรหัสข้อความ AES
    private String decrypt(String encryptedMessage, SecretKeySpec SecretKeyAESReceive) throws Exception {
    	
    	logger.Info("Before decrypt message with AES         ::: " + encryptedMessage);
    	
        cipherAES.init(Cipher.DECRYPT_MODE, SecretKeyAESReceive, ivspec);
    	byte[] clearbyte = Base64.decodeBase64(encryptedMessage);
        byte[] original = cipherAES.doFinal(clearbyte);
        
        logger.Info("After decrypt message with AES          ::: " + new String(original));
        
        return new String(original,"UTF8");
    }
    
    // นำรหัส AES ที่ถูกเข้ารหัส RSA นำมาถอดรหัส RSA ออก
    private SecretKeySpec decryptKeyAES(byte[] AESKeyReceive) throws Exception {
    	cipherRSA.init(Cipher.DECRYPT_MODE, privateKeyRSA);
    	decryptedRawKey = cipherRSA.doFinal(AESKeyReceive);
    	
    	logger.Info("Decrypt key AES with RSA                ::: " + new SecretKeySpec(decryptedRawKey, "AES"));
    	
    	return new SecretKeySpec(decryptedRawKey, "AES");
    }
    
    // นำรหัส AES เข้า รหัส RSA
    private void encryptKeyAES() throws Exception {
    	convertInterlocutorPublicKey(getinterlocutorPublicKey);
    	cipherRSA.init(Cipher.ENCRYPT_MODE, interlocutorPublicKey);
        encryptedAesKey = cipherRSA.doFinal(rawKey);
        
        logger.Info("Encrypt key AES with RSA                ::: " + encryptedAesKey);
    }
    
    // นำรหัส HMAC เข้า รหัส RSA
    private void encryptKeyHMAC() throws Exception {
    	convertInterlocutorPublicKey(getinterlocutorPublicKey);
    	cipherRSA.init(Cipher.ENCRYPT_MODE, interlocutorPublicKey);
    	encryptedHMACKey = cipherRSA.doFinal(HMACKey.getBytes());
    	
    	logger.Info("Encrypt key HMAC with RSA               ::: " + encryptedHMACKey);
    }
    
    // แปลง Public Key ของคู่สนทนาเป็น Key
    private void convertInterlocutorPublicKey(byte[] interlocutorPublicKey) throws Exception {
    	keyFactory = KeyFactory.getInstance("RSA");
	    keySpec = new X509EncodedKeySpec(interlocutorPublicKey);
	    this.interlocutorPublicKey = keyFactory.generatePublic(keySpec);
	    
	    logger.Info("Convert Interlocutor   key[byte] to Key ::: " + interlocutorPublicKey);
    }
    
    // นำรหัส HMAC ที่ถูกเข้ารหัส RSA นำมาถอดรหัส RSA ออก
    private String decryptKeyIntegrity(byte[] HMACKeyReceive) throws Exception {
    	cipherRSA.init(Cipher.DECRYPT_MODE, privateKeyRSA);
    	decryptedRawKey = cipherRSA.doFinal(HMACKeyReceive);
    	
    	logger.Info("Decrypt key HMAC with RSA               ::: " + decryptedRawKey);
    	
    	return new String(decryptedRawKey);
    }
    
    //================================================================================
    // Accessors
    //================================================================================
    
    // รับค่า PublicKey ของตัวเอง
    public byte[] getMyPublicKey() {
    	logger.Info("Get my public key[byte]                 ::: " + publicKeyRSA.getEncoded());
    	
        return publicKeyRSA.getEncoded();
    }

    // เก็บค่า PublicKey ของคู่สนทนา
    public void setInterlocutorPublicKey(byte[] interlocutorPublicKey) {
    	logger.Info("Set interlocutor public key[byte]       ::: " + interlocutorPublicKey);
    	
        this.getinterlocutorPublicKey = interlocutorPublicKey;
    }
    
}
